#!/bin/bash

# 存储性能测试脚本
# 参数包括: 读写模式(随机/顺序)、IO大小(块大小)、读写比例和并发度(线程数)

# 默认参数值
IO_PATTERN="rand"    # 默认随机读写, 可选值: rand(随机), seq(顺序)
BLOCK_SIZE="4k"      # 默认块大小4KB
RW_RATIO="70"        # 默认读写比例70%(读)/30%(写)
THREADS="4"          # 默认4个线程
TEST_FILE="/tmp/fio_test_file"
TEST_SIZE="1G"       # 默认测试文件大小
RUNTIME="60"         # 默认运行时间(秒)
TIMESTAMP=$(date +"%Y%m%d_%H%M%S")
LOG_FILE="storage_test_${TIMESTAMP}.txt"  # 默认日志文件名，使用.txt扩展名

# 函数: 显示帮助信息
show_help() {
    echo "用法: $0 [选项]"
    echo "选项:"
    echo "  -p, --pattern PATTERN    IO模式: rand(随机) 或 seq(顺序), 默认: $IO_PATTERN"
    echo "  -b, --blocksize SIZE     IO块大小, 例如: 4k, 16k, 1m, 默认: $BLOCK_SIZE"
    echo "  -r, --rwratio RATIO      读写比例(读百分比), 0-100, 默认: $RW_RATIO"
    echo "  -t, --threads NUM        并发线程数, 默认: $THREADS"
    echo "  -f, --file PATH          测试文件路径, 默认: $TEST_FILE"
    echo "  -s, --size SIZE          测试文件大小, 默认: $TEST_SIZE"
    echo "  -d, --runtime DURATION   测试运行时间(秒), 默认: $RUNTIME"
    echo "  -h, --help               显示帮助信息"
    exit 0
}

# 解析命令行参数
while [[ $# -gt 0 ]]; do
    case $1 in
        -p|--pattern)
            IO_PATTERN="$2"
            if [[ "$IO_PATTERN" != "rand" && "$IO_PATTERN" != "seq" ]]; then
                echo "错误: IO模式必须是 'rand' 或 'seq'"
                exit 1
            fi
            shift 2
            ;;
        -b|--blocksize)
            BLOCK_SIZE="$2"
            shift 2
            ;;
        -r|--rwratio)
            RW_RATIO="$2"
            if ! [[ "$RW_RATIO" =~ ^[0-9]+$ ]] || [ "$RW_RATIO" -gt 100 ]; then
                echo "错误: 读写比例必须是0-100之间的整数"
                exit 1
            fi
            shift 2
            ;;
        -t|--threads)
            THREADS="$2"
            if ! [[ "$THREADS" =~ ^[0-9]+$ ]] || [ "$THREADS" -lt 1 ]; then
                echo "错误: 线程数必须是正整数"
                exit 1
            fi
            shift 2
            ;;
        -f|--file)
            TEST_FILE="$2"
            shift 2
            ;;
        -s|--size)
            TEST_SIZE="$2"
            shift 2
            ;;
        -d|--runtime)
            RUNTIME="$2"
            shift 2
            ;;
        -h|--help)
            show_help
            ;;
        *)
            echo "未知选项: $1"
            show_help
            ;;
    esac
done

# 更新日志文件名，包含测试参数
LOG_FILE="storage_test_${IO_PATTERN}_${BLOCK_SIZE}_r${RW_RATIO}_t${THREADS}_${TIMESTAMP}.txt"

# 检查FIO是否已安装
if ! command -v fio &> /dev/null; then
    echo "错误: fio 未安装，请先安装 Flexible I/O Tester"
    echo "  Debian/Ubuntu: sudo apt-get install fio"
    echo "  RHEL/CentOS:   sudo yum install fio"
    exit 1
fi

# 设置FIO参数
TEST_SUMMARY="=== 存储性能测试 ===
IO模式: ${IO_PATTERN} (rand=随机, seq=顺序)
块大小: $BLOCK_SIZE
读写比例: ${RW_RATIO}% 读 / $((100-RW_RATIO))% 写
并发线程数: $THREADS
测试文件: $TEST_FILE
文件大小: $TEST_SIZE
运行时间: ${RUNTIME}秒
日志文件: $LOG_FILE
===================="

echo "$TEST_SUMMARY"

# 构建FIO命令
RW_TYPE=""
if [ "$IO_PATTERN" = "rand" ]; then
    if [ "$RW_RATIO" = "100" ]; then
        RW_TYPE="randread"
    elif [ "$RW_RATIO" = "0" ]; then
        RW_TYPE="randwrite"
    else
        RW_TYPE="randrw"
    fi
else
    if [ "$RW_RATIO" = "100" ]; then
        RW_TYPE="read"
    elif [ "$RW_RATIO" = "0" ]; then
        RW_TYPE="write"
    else
        RW_TYPE="rw"
    fi
fi

# 创建临时FIO配置文件
CONFIG_FILE="/tmp/fio_test_$$.ini"
cat > "$CONFIG_FILE" << EOF
[global]
name=storage_test
filename=${TEST_FILE}
size=${TEST_SIZE}
direct=1
verify=0
ioengine=libaio
iodepth=64
runtime=${RUNTIME}
numjobs=${THREADS}
time_based=1
group_reporting=1

[job]
rw=${RW_TYPE}
rwmixread=${RW_RATIO}
bs=${BLOCK_SIZE}
EOF

echo "开始测试..."
echo "测试将持续 ${RUNTIME} 秒..."
echo "测试结果将保存到: $LOG_FILE"

# 创建临时输出文件
FIO_OUTPUT="/tmp/fio_output_$$.txt"

# 执行FIO测试并捕获输出
fio "$CONFIG_FILE" | tee "$FIO_OUTPUT"

# 解析FIO输出并提取关键指标
analyze_fio_output() {
    local fio_file="$1"
    local log_output=""
    
    # 添加测试摘要
    log_output+="$TEST_SUMMARY\n\n"
    
    # 1. 提取读写摘要
    local read_iops=$(grep -o "read: IOPS=[0-9.]*[kKmM]*" "$fio_file" | cut -d= -f2)
    local read_bw=$(grep -o "read:.*BW=[0-9.]*[kKmMgG]*iB/s" "$fio_file" | grep -o "BW=[0-9.]*[kKmMgG]*iB/s" | cut -d= -f2)
    local read_io=$(grep -A1 "Run status group" "$fio_file" | grep "READ" | grep -o "io=[0-9.]*[kKmMgG]*iB" | cut -d= -f2)
    local write_iops=$(grep -o "write: IOPS=[0-9.]*[kKmM]*" "$fio_file" | cut -d= -f2)
    local write_bw=$(grep -o "write:.*BW=[0-9.]*[kKmMgG]*iB/s" "$fio_file" | grep -o "BW=[0-9.]*[kKmMgG]*iB/s" | cut -d= -f2)
    local write_io=$(grep -A2 "Run status group" "$fio_file" | grep "WRITE" | grep -o "io=[0-9.]*[kKmMgG]*iB" | cut -d= -f2)
    
    log_output+="------------------------------------------\n"
    log_output+="READ/WRITE SUMMARY\n"
    log_output+="------------------------------------------\n"
    log_output+="读取性能:\n"
    log_output+="  IOPS: $read_iops (每秒IO操作数)\n"
    log_output+="  带宽: $read_bw\n"
    log_output+="  总读取量: $read_io\n\n"
    log_output+="写入性能:\n"
    log_output+="  IOPS: $write_iops (每秒IO操作数)\n"
    log_output+="  带宽: $write_bw\n"
    log_output+="  总写入量: $write_io\n\n"
    
    # 2. 提取延迟分布 - 只保留延迟百分比分布行
    local lat_usec_pct=$(grep "lat (usec)" "$fio_file" | grep -v "min=\|max=\|avg=")
    local lat_msec_pct=$(grep "lat (msec)" "$fio_file" | grep -v "min=\|max=\|avg=")
    
    log_output+="------------------------------------------\n"
    log_output+="GLOBAL LATENCY DISTRIBUTION\n"
    log_output+="------------------------------------------\n"
    if [ -n "$lat_usec_pct" ]; then
        log_output+="$lat_usec_pct\n"
    fi
    if [ -n "$lat_msec_pct" ]; then
        log_output+="$lat_msec_pct\n"
    fi
    log_output+="\n"
    
    # 3. 提取系统资源和队列深度 - 简化输出
    local cpu_usr=$(grep "cpu " "$fio_file" | grep -o "usr=[0-9.]*%" | cut -d= -f2)
    local cpu_sys=$(grep "cpu " "$fio_file" | grep -o "sys=[0-9.]*%" | cut -d= -f2)
    local io_depths=$(grep "IO depths" "$fio_file" | grep -A1 "IO depths" | head -1)
    
    log_output+="------------------------------------------\n"
    log_output+="SYSTEM & QUEUE STATS\n"
    log_output+="------------------------------------------\n"
    log_output+="CPU使用率:\n"
    log_output+="  用户态: $cpu_usr\n"
    log_output+="  系统态: $cpu_sys\n\n"
    log_output+="IO队列深度:\n"
    log_output+="  $io_depths\n\n"
    
    # 4. 提取磁盘统计和运行状态
    local dm_util=$(grep -A2 "Disk stats" "$fio_file" | grep "dm-0" | grep -o "util=[0-9.]*%" | cut -d= -f2)
    local nvme_util=$(grep -A2 "Disk stats" "$fio_file" | grep "nvme" | grep -o "util=[0-9.]*%" | cut -d= -f2)
    
    log_output+="------------------------------------------\n"
    log_output+="DISK STATS & RUN STATUS\n"
    log_output+="------------------------------------------\n"
    log_output+="最终状态摘要:\n"
    log_output+="  READ: 带宽=$read_bw, 总量=$read_io\n"
    log_output+="  WRITE: 带宽=$write_bw, 总量=$write_io\n\n"
    log_output+="磁盘利用率:\n"
    if [ -n "$dm_util" ]; then
        log_output+="  逻辑卷(dm-0): $dm_util\n"
    fi
    if [ -n "$nvme_util" ]; then
        log_output+="  物理设备(nvme): $nvme_util\n"
    fi
    log_output+="\n"
    
    # 5. 原始FIO输出
    log_output+="------------------------------------------\n"
    log_output+="ORIGINAL FIO OUTPUT\n"
    log_output+="------------------------------------------\n"
    log_output+="$(cat "$fio_file")\n"
    
    echo -e "$log_output"
}

# 分析FIO输出
ANALYSIS=$(analyze_fio_output "$FIO_OUTPUT")

# 输出到日志文件
echo -e "$ANALYSIS" > "$LOG_FILE"
echo "测试结果已保存到: $LOG_FILE"

# 清理临时文件
rm -f "$CONFIG_FILE" "$FIO_OUTPUT"
echo "测试完成"

# 提示是否需要删除测试文件
read -p "是否删除测试文件 ${TEST_FILE}? (y/n): " delete_file
if [[ "$delete_file" == "y" || "$delete_file" == "Y" ]]; then
    rm -f "$TEST_FILE"
    echo "测试文件已删除"
fi 
